{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "nerfies_2d_experiments.ipynb",
      "provenance": [
        {
          "file_id": "1wdNjll4hs3b8yy7O0IKpXSXrXAyOVEaY",
          "timestamp": 1631790666234
        },
        {
          "file_id": "1Zv_LdAL82dfOruTngKaSLViawaa3SkL6",
          "timestamp": 1615009722746
        }
      ],
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      }
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "N_PcTdlrG3nu"
      },
      "source": [
        "# @title Imports\n",
        "\n",
        "from dataclasses import dataclass\n",
        "from pprint import pprint\n",
        "from typing import Any, List, Callable, Dict, Sequence, Optional, Tuple\n",
        "from io import BytesIO\n",
        "from IPython.display import display, HTML\n",
        "from base64 import b64encode\n",
        "import PIL\n",
        "import IPython\n",
        "import tempfile\n",
        "import imageio\n",
        "\n",
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "import tensorflow as tf\n",
        "\n",
        "import jax\n",
        "from jax.config import config as jax_config\n",
        "import jax.numpy as jnp\n",
        "from jax import grad, jit, vmap\n",
        "from jax import random\n",
        "\n",
        "import flax\n",
        "import flax.linen as nn\n",
        "# from flax import nn\n",
        "from flax import jax_utils\n",
        "from flax import optim\n",
        "from flax.metrics import tensorboard\n",
        "from flax.training import checkpoints\n",
        "\n",
        "from absl import logging\n",
        "\n",
        "# Monkey patch logging.\n",
        "def myprint(msg, *args, **kwargs):\n",
        " print(msg % args)\n",
        "\n",
        "logging.info = myprint \n",
        "logging.warn = myprint"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u4kI6OZ0MwPP"
      },
      "source": [
        "## Utility Functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bph93J0PMyeI",
        "cellView": "form"
      },
      "source": [
        "# @title Dataset Utilities\n",
        "\n",
        "def prepare_tf_data(xs):\n",
        "  \"\"\"Convert a input batch from tf Tensors to numpy arrays.\"\"\"\n",
        "  local_device_count = jax.local_device_count()\n",
        "  def _prepare(x):\n",
        "    # Use _numpy() for zero-copy conversion between TF and NumPy.\n",
        "    x = x._numpy()  # pylint: disable=protected-access\n",
        "\n",
        "    # reshape (host_batch_size, height, width, 3) to\n",
        "    # (local_devices, device_batch_size, height, width, 3)\n",
        "    return x.reshape((local_device_count, -1) + x.shape[1:])\n",
        "\n",
        "  return jax.tree_map(_prepare, xs)\n",
        "\n",
        "\n",
        "def iterator_from_dataset(dataset: tf.data.Dataset,\n",
        "                          batch_size: int,\n",
        "                          repeat: bool = True,\n",
        "                          prefetch_size: int = 0,\n",
        "                          devices: Optional[Sequence[Any]] = None):\n",
        "  \"\"\"Create a data iterator that returns JAX arrays from a TF dataset.\n",
        "\n",
        "  Args:\n",
        "    dataset: the dataset to iterate over.\n",
        "    batch_size: the batch sizes the iterator should return.\n",
        "    repeat: whether the iterator should repeat the dataset.\n",
        "    prefetch_size: the number of batches to prefetch to device.\n",
        "    devices: the devices to prefetch to.\n",
        "\n",
        "  Returns:\n",
        "    An iterator that returns data batches.\n",
        "  \"\"\"\n",
        "  if repeat:\n",
        "    dataset = dataset.repeat()\n",
        "\n",
        "  if batch_size > 0:\n",
        "    dataset = dataset.batch(batch_size)\n",
        "    it = map(prepare_tf_data, dataset)\n",
        "  else:\n",
        "    it = map(prepare_tf_data_unbatched, dataset)\n",
        "\n",
        "  if prefetch_size > 0:\n",
        "    it = jax_utils.prefetch_to_device(it, prefetch_size, devices)\n",
        "\n",
        "  return it\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ux0cH2TpQyEN"
      },
      "source": [
        "## Create Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ry0smCvRyttR"
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "from matplotlib import cm\n",
        "import matplotlib as mpl\n",
        "\n",
        "\n",
        "def fig_to_array(fig, height, width, dpi=100):\n",
        "  out_array = np.zeros((height, width, 4), dtype=np.uint8)\n",
        "  fig.set_size_inches((width / dpi, height / dpi))\n",
        "  fig.set_dpi(dpi)\n",
        "  for ax in fig.axes:\n",
        "    ax.margins(0, 0)\n",
        "    ax.axis('off')\n",
        "    ax.get_xaxis().set_visible(False)\n",
        "    ax.get_yaxis().set_visible(False)\n",
        "\n",
        "  # If we haven't already shown or saved the plot, then we need to\n",
        "  # draw the figure first...\n",
        "  fig.tight_layout(pad=0)\n",
        "  fig.canvas.draw()\n",
        "\n",
        "  # Now we can save it to a numpy array.\n",
        "  data = np.frombuffer(fig.canvas.tostring_argb(), dtype=np.uint8)\n",
        "  data = data.reshape(fig.canvas.get_width_height()[::-1] + (4,))\n",
        "  plt.close()\n",
        "\n",
        "  np.copyto(out_array, data)\n",
        "  out_array = np.roll(out_array, -1, 2)\n",
        "  return out_array\n",
        "\n",
        "\n",
        "def add_colorwheel(fig, cmap, lim, label):\n",
        "  display_axes = fig.add_axes([0.1,0.1,0.8,0.8], projection='polar', label=label)\n",
        "  display_axes._direction = 2*np.pi ## This is a nasty hack - using the hidden field to \n",
        "                                    ## multiply the values such that 1 become 2*pi\n",
        "                                    ## this field is supposed to take values 1 or -1 only!!\n",
        "  \n",
        "  norm = mpl.colors.Normalize(0.0, 2*np.pi)\n",
        "  \n",
        "  # Plot the colorbar onto the polar axis\n",
        "  # note - use orientation horizontal so that the gradient goes around\n",
        "  # the wheel rather than centre out\n",
        "  quant_steps = 2056\n",
        "  cb = mpl.colorbar.ColorbarBase(display_axes, cmap=cm.get_cmap(cmap,quant_steps),\n",
        "                                norm=norm,\n",
        "                                orientation='horizontal')\n",
        "  \n",
        "  # aesthetics - get rid of border and axis labels                                   \n",
        "  cb.outline.set_visible(False)                                 \n",
        "  display_axes.set_axis_off()\n",
        "  display_axes.set_rlim(lim)\n",
        "\n",
        "  display_axes.margins(0, 0)\n",
        "  display_axes.axis('off')\n",
        "  display_axes.get_xaxis().set_visible(False)\n",
        "  display_axes.get_yaxis().set_visible(False)\n",
        "  display_axes.set_aspect(1)\n",
        "\n",
        "fig = plt.figure()\n",
        "add_colorwheel(fig, 'tab20b', (-1.1, 1.0), 'a')\n",
        "add_colorwheel(fig, 'twilight', (-2.0, 1), 'b')\n",
        "add_colorwheel(fig, 'hsv', (-5.0, 1), 'c')\n",
        "\n",
        "colorwheel = fig_to_array(fig, 400, 400)\n",
        "mediapy.show_image(colorwheel)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v-_vZ8cfFtcv"
      },
      "source": [
        "# import matplotlib\n",
        "import scipy\n",
        "from PIL import Image, ImageDraw, ImageFont, ImageOps\n",
        "\n",
        "def barron_colormap(n):\n",
        "  curve = lambda x, s, t: np.where(x<t, (t*x)/(x+s*(t-x)+1e-10), ((1-t)*(x-1))/(1-x-s*(t-x)+1e-10)+1)\n",
        "  t = curve(jnp.linspace(0, 1, n), 1.5, 0.5)\n",
        "  colors = plt.get_cmap('rainbow')(t)[:,:3]\n",
        "  colors = 0.85 * (1 - (1-colors) / jnp.sqrt(jnp.sum((1-colors)**2, 1, keepdims=True)))\n",
        "  return colors\n",
        "\n",
        "\n",
        "def make_doodle(height, width, tile_size=5):\n",
        "  h = height + (height - 1) % 2\n",
        "  w = width + (width - 1) % 2\n",
        "  colors = np.array(barron_colormap(h * w))\n",
        "  colors[::2, :] = 0.0\n",
        "  colors = colors.reshape((h, w, 3))\n",
        "  colors = colors[:h, :w]\n",
        "  colors = np.repeat(colors, tile_size, axis=0)\n",
        "  colors = np.repeat(colors, tile_size, axis=1)\n",
        "  return np.asarray(colors)\n",
        "\n",
        "\n",
        "# doodle = make_doodle(20, 20, 4)\n",
        "# files = %upload_files\n",
        "# doodles = []\n",
        "# for f in files.values():\n",
        "#   doodles.append(imageio.imread(f))\n",
        "doodles = [image_utils.load_image('')]\n",
        "mediapy.show_images(doodles)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wixxRenKggPe"
      },
      "source": [
        "from PIL import ImageChops, ImageFilter\n",
        "import math\n",
        "import skimage\n",
        "from skimage.transform import swirl\n",
        "\n",
        "# def random_image(rng, shape, doodles):\n",
        "#   h, w = shape[:2]\n",
        "#   num_doodles = len(doodles)\n",
        "#   image = Image.new('RGBA', shape[::-1], color=(127, 127, 127))\n",
        "#   for i, doodle in enumerate(doodles):\n",
        "#     dh, dw = doodle.shape[:2]\n",
        "#     rng, k1, k2, k3 = random.split(rng, 4)\n",
        "#     angle = random.uniform(k1, (1,), minval=-180, maxval=180)[0]\n",
        "#     y = random.randint(k2, (1,), minval=0, maxval=h-int(math.sqrt(dh**2+dw**2)))\n",
        "#     x = random.randint(k3, (1,), minval=0, maxval=w-int(math.sqrt(dh**2+dw**2)))\n",
        "#     doodle = Image.fromarray(image_utils.image_to_uint8(doodle)).convert('RGBA')\n",
        "#     doodle = doodle.rotate(angle, expand=True, resample=Image.BICUBIC)\n",
        "#     image.paste(doodle, (x, y), mask=doodle)\n",
        "#   return np.asarray(image.convert('RGB')).astype(np.float32) / 255.0\n",
        "\n",
        "\n",
        "def crop_circle(array, blur_radius=0, offset=0):\n",
        "  pil_img = Image.fromarray(image_utils.image_to_uint8(array))\n",
        "  offset = blur_radius * 2 + offset\n",
        "  mask = Image.new(\"L\", pil_img.size, 0)\n",
        "  draw = ImageDraw.Draw(mask)\n",
        "  draw.ellipse((offset, offset, pil_img.size[0] - offset, pil_img.size[1] - offset), fill=255)\n",
        "  mask = mask.filter(ImageFilter.GaussianBlur(blur_radius))\n",
        "\n",
        "  result = pil_img.copy()\n",
        "  result.putalpha(mask)\n",
        "\n",
        "  return np.asarray(result)\n",
        "\n",
        "\n",
        "def random_image(rng, shape, frame, doodle, shift=10):\n",
        "  rng, k1, k2, k3, k4 = random.split(rng, 5)\n",
        "\n",
        "  h, w = shape\n",
        "  dh, dw = doodle.shape[:2]\n",
        "\n",
        "  rng, swirl_key1, swirl_key2, swirl_key3 = random.split(rng, 4)\n",
        "  swirl_angle = random.uniform(swirl_key1, (1,), minval=-0, maxval=0)[0]\n",
        "  swirl_strength = random.uniform(swirl_key2, (1,), minval=-6, maxval=6)[0]\n",
        "  swirl_center_dx, swirl_center_dy = random.uniform(swirl_key3, (2,), minval=-20, maxval=20)\n",
        "  doodle = swirl(doodle, \n",
        "                 rotation=swirl_angle, \n",
        "                 strength=swirl_strength, \n",
        "                 radius=200,\n",
        "                 center=(dh//2+swirl_center_dx, dw//2+swirl_center_dy))\n",
        "  doodle = crop_circle(doodle)\n",
        "\n",
        "  angle = random.uniform(k1, (1,), minval=-45, maxval=45)[0]\n",
        "  doodle = Image.fromarray(image_utils.image_to_uint8(doodle))\n",
        "  doodle = doodle.copy()\n",
        "  # doodle = doodle.rotate(angle, expand=True, resample=Image.BICUBIC)\n",
        "  frame = Image.fromarray(image_utils.image_to_uint8(frame))\n",
        "  doodle_x = int(frame.width / 2 - doodle.width / 2)\n",
        "  doodle_y = int(frame.height / 2 - doodle.height / 2)\n",
        "  frame.paste(doodle, (doodle_x, doodle_y), mask=doodle)\n",
        "\n",
        "  image = Image.new('RGBA', shape[::-1], color=(50, 50, 50))\n",
        "  angle = random.uniform(k4, (1,), minval=-180, maxval=180)[0]\n",
        "  cx = int(image.width / 2 - frame.width / 2)\n",
        "  cy = int(image.height / 2 - frame.height / 2)\n",
        "  y = random.randint(k2, (1,), minval=cx-shift, maxval=cx+shift)\n",
        "  x = random.randint(k3, (1,), minval=cy-shift, maxval=cy+shift)\n",
        "\n",
        "  frame = frame.rotate(angle, expand=False, resample=Image.BICUBIC)\n",
        "  image.paste(frame, (x, y), mask=frame)\n",
        "\n",
        "  return image_utils.image_to_float32(np.asarray(image))\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "num_images = 20\n",
        "image_shape = (400, 400)\n",
        "doodle_scale = 1/7\n",
        "# doodles = [skimage.data.astronaut()]\n",
        "rescaled_doodles = [image_utils.rescale_image(d, doodle_scale) for d in doodles]\n",
        "imedia.show_images(rescaled_doodles)\n",
        "shift = 40\n",
        "seed = 6\n",
        "rng = random.PRNGKey(seed)\n",
        "\n",
        "images = []\n",
        "for i in range(num_images):\n",
        "  rng, key = random.split(rng)\n",
        "  images.append(random_image(key, image_shape, colorwheel, rescaled_doodles[0], shift=shift))\n",
        "\n",
        "imedia.show_images(images, border=True, columns=5)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lEasEtAzyQb6"
      },
      "source": [
        "for image_idx, image in enumerate(images):\n",
        "  path = gpath.GPath('', f'{image_idx:02d}.png')\n",
        "  path.parent.mkdir(exist_ok=True, parents=True)\n",
        "  image_utils.save_image(path, image_utils.image_to_uint8(image))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tLQuD7XkX3QV"
      },
      "source": [
        "def image_to_data(image, idx, scale=1.0):\n",
        "  coords = jnp.stack(jnp.meshgrid(\n",
        "    jnp.linspace(-scale, scale, image.shape[1]),\n",
        "    jnp.linspace(-scale, scale, image.shape[0]),\n",
        "  ), axis=-1)\n",
        "  return {\n",
        "      'coords': coords,\n",
        "      'colors': image[..., :3],\n",
        "      'ids': jnp.full_like(image[:, :, 0], idx, dtype=jnp.uint32),\n",
        "  }\n",
        "\n",
        "data_items = [image_to_data(image, i) for i, image in enumerate(images)]\n",
        "data = jax.tree_multimap(lambda *x: jnp.stack(x), *data_items)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "idZCoRq5EqlJ"
      },
      "source": [
        "imedia.show_images([data['coords'][0, ..., 0], data['coords'][0, ..., 1]], cmap='coolwarm')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "38NCFU1JM6_X"
      },
      "source": [
        "## Train"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5wL17_AHUD3m"
      },
      "source": [
        "# @title MLP\n",
        "\n",
        "class MLP(nn.Module):\n",
        "  \"\"\"Basic MLP class with hidden layers and an output layers.\"\"\"\n",
        "  depth: int\n",
        "  width: int\n",
        "  hidden_init: Any\n",
        "  hidden_activation: Any = nn.relu\n",
        "  output_init: Optional[Any] = None\n",
        "  output_channels: int = 0\n",
        "  output_activation: Optional[Any] = lambda x: x\n",
        "  use_bias: bool = True\n",
        "  skips: Tuple[int] = tuple()\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x):\n",
        "    inputs = x\n",
        "    for i in range(self.depth):\n",
        "      layer = nn.Dense(\n",
        "          self.width,\n",
        "          use_bias=self.use_bias,\n",
        "          kernel_init=self.hidden_init,\n",
        "          name=f'hidden_{i}')\n",
        "      if i in self.skips:\n",
        "        x = jnp.concatenate([x, inputs], axis=-1)\n",
        "      x = layer(x)\n",
        "      x = self.hidden_activation(x)\n",
        "\n",
        "    if self.output_channels > 0:\n",
        "      logit_layer = nn.Dense(\n",
        "          self.output_channels,\n",
        "          use_bias=self.use_bias,\n",
        "          kernel_init=self.output_init,\n",
        "          name='logit')\n",
        "      x = logit_layer(x)\n",
        "      if self.output_activation is not None:\n",
        "        x = self.output_activation(x)\n",
        "\n",
        "    return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XoI2JfIeUBIW",
        "cellView": "form"
      },
      "source": [
        "# @title SinusoidalEncoder\n",
        " \n",
        "\n",
        "class SinusoidalEncoder(nn.Module):\n",
        "  \"\"\"A vectorized sinusoidal encoding.\n",
        "\n",
        "  Attributes:\n",
        "    num_freqs: the number of frequency bands in the encoding.\n",
        "    max_freq_log2: the log (base 2) of the maximum frequency.\n",
        "    scale: a scaling factor for the positional encoding.\n",
        "    use_identity: if True use the identity encoding as well.\n",
        "  \"\"\"\n",
        "  num_freqs: int\n",
        "  min_freq_log2: int = 0\n",
        "  max_freq_log2: Optional[int] = None\n",
        "  scale: float = 1.0\n",
        "  use_identity: bool = True\n",
        "\n",
        "  def setup(self):\n",
        "    if self.max_freq_log2 is None:\n",
        "      max_freq_log2 = self.num_freqs - 1.0\n",
        "    else:\n",
        "      max_freq_log2 = self.max_freq_log2\n",
        "    self.freq_bands = 2.0**jnp.linspace(self.min_freq_log2, max_freq_log2, int(self.num_freqs))\n",
        "\n",
        "    # (F, 1).\n",
        "    self.freqs = jnp.reshape(self.freq_bands, (self.num_freqs, 1))\n",
        "\n",
        "  def __call__(self, x, alpha: Optional[float] = None):\n",
        "    \"\"\"A vectorized sinusoidal encoding.\n",
        "\n",
        "    Args:\n",
        "      x: the input features to encode.\n",
        "      alpha: a dummy argument for API compatibility.\n",
        "\n",
        "    Returns:\n",
        "      A tensor containing the encoded features.\n",
        "    \"\"\"\n",
        "    if self.num_freqs == 0:\n",
        "      return x\n",
        "\n",
        "    x_expanded = jnp.expand_dims(x, axis=-2)  # (1, C).\n",
        "    # Will be broadcasted to shape (F, C).\n",
        "    angles = self.scale * x_expanded * self.freqs\n",
        "\n",
        "    # The shape of the features is (F, 2, C) so that when we reshape it\n",
        "    # it matches the ordering of the original NeRF code.\n",
        "    # Vectorize the computation of the high-frequency (sin, cos) terms.\n",
        "    # We use the trigonometric identity: cos(x) = sin(x + pi/2)\n",
        "    features = jnp.stack((angles, angles + jnp.pi / 2), axis=-2)\n",
        "    features = features.flatten()\n",
        "    features = jnp.sin(features)\n",
        "\n",
        "    # Prepend the original signal for the identity.\n",
        "    if self.use_identity:\n",
        "      features = jnp.concatenate([x, features], axis=-1)\n",
        "    return features\n",
        "\n",
        "\n",
        "class AnnealedSinusoidalEncoder(nn.Module):\n",
        "  \"\"\"An annealed sinusoidal encoding.\"\"\"\n",
        "  num_freqs: int\n",
        "  min_freq_log2: int = 0\n",
        "  max_freq_log2: Optional[int] = None\n",
        "  scale: float = 1.0\n",
        "  use_identity: bool = False\n",
        "\n",
        "  @nn.compact\n",
        "  def __call__(self, x, alpha):\n",
        "    if alpha is None:\n",
        "      raise ValueError('alpha must be specified.')\n",
        "    if self.num_freqs == 0:\n",
        "      return x\n",
        "\n",
        "    num_channels = x.shape[-1]\n",
        "\n",
        "    base_encoder = SinusoidalEncoder(\n",
        "        num_freqs=self.num_freqs,\n",
        "        min_freq_log2=self.min_freq_log2,\n",
        "        max_freq_log2=self.max_freq_log2,\n",
        "        scale=self.scale,\n",
        "        use_identity=self.use_identity)\n",
        "    features = base_encoder(x)\n",
        "\n",
        "    if self.use_identity:\n",
        "      identity, features = jnp.split(features, (x.shape[-1],), axis=-1)\n",
        "\n",
        "    # Apply the window by broadcasting to save on memory.\n",
        "    features = jnp.reshape(features, (-1, 2, num_channels))\n",
        "    window = self.cosine_easing_window(\n",
        "        self.min_freq_log2, self.max_freq_log2, self.num_freqs, alpha)\n",
        "    window = jnp.reshape(window, (-1, 1, 1))\n",
        "    features = window * features\n",
        "\n",
        "    if self.use_identity:\n",
        "      return jnp.concatenate([\n",
        "          identity,\n",
        "          features.flatten(),\n",
        "      ], axis=-1)\n",
        "    else:\n",
        "      return features.flatten()\n",
        "\n",
        "  @classmethod\n",
        "  def cosine_easing_window(cls, min_freq_log2, max_freq_log2, num_bands, alpha):\n",
        "    \"\"\"Eases in each frequency one by one with a cosine.\n",
        "\n",
        "    This is equivalent to taking a Tukey window and sliding it to the right\n",
        "    along the frequency spectrum.\n",
        "\n",
        "    Args:\n",
        "      num_freqs: the number of frequencies.\n",
        "      alpha: will ease in each frequency as alpha goes from 0.0 to num_freqs.\n",
        "\n",
        "    Returns:\n",
        "      A 1-d numpy array with num_sample elements containing the window.\n",
        "    \"\"\"\n",
        "    bands = jnp.linspace(min_freq_log2,  max_freq_log2, num_bands)\n",
        "    x = jnp.clip(alpha - bands, 0.0, 1.0)\n",
        "    return 0.5 * (1 + jnp.cos(jnp.pi * x + jnp.pi))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YR8smVFZKRif"
      },
      "source": [
        "devices = jax.devices()\n",
        "rng = random.PRNGKey(1)\n",
        "rng, key = random.split(rng, 2)\n",
        "\n",
        "template_min_freq = -2.0\n",
        "template_max_freq = 8.0\n",
        "template_num_freqs = int(template_max_freq - template_min_freq + 1)\n",
        "scale = 1.0\n",
        "deform_type = 'translation'\n",
        "\n",
        "# Anneal.\n",
        "# use_anneal = True\n",
        "# min_freq = -2.0\n",
        "# max_freq = 4.0\n",
        "\n",
        "# No anneal.\n",
        "use_anneal = False\n",
        "min_freq = -2.0\n",
        "max_freq = -2.0\n",
        "\n",
        "# Common.\n",
        "num_freqs = int(max_freq - min_freq + 1)\n",
        "\n",
        "\n",
        "VModel = nn.vmap(Model, \n",
        "                 in_axes=(0, 0, None, None),\n",
        "                 variable_axes={'params': None},  \n",
        "                 split_rngs={'params': False})\n",
        "model = VModel(num_glo_embeddings=len(images),\n",
        "               deform_type=deform_type,\n",
        "               deform_min_freq_log2=min_freq,\n",
        "               deform_max_freq_log2=max_freq,\n",
        "               deform_num_freqs=num_freqs,\n",
        "               template_num_freqs=template_num_freqs,\n",
        "               template_min_freq_log2=template_min_freq,\n",
        "               template_max_freq_log2=template_max_freq,\n",
        "               scale=scale)\n",
        "init_coords = random.normal(key, (1024, 2))\n",
        "init_glo_ids = jnp.zeros((1024,), jnp.uint32)\n",
        "init_params = model.init(key, init_coords, init_glo_ids, 0.0, True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O1Xy1qY9JKkE"
      },
      "source": [
        "# Setup dataset\n",
        "def prepare_data(xs):\n",
        "  \"\"\"Convert a input batch from tf Tensors to numpy arrays.\"\"\"\n",
        "  local_device_count = jax.local_device_count()\n",
        "  def _prepare(x):\n",
        "    # reshape (host_batch_size, height, width, 3) to\n",
        "    # (local_devices, device_batch_size, height, width, 3)\n",
        "    x = x.reshape((local_device_count, -1) + x.shape[1:])\n",
        "    return jax.api.device_put_sharded(list(x), jax.local_devices())\n",
        "\n",
        "  return jax.tree_map(_prepare, xs)\n",
        "\n",
        "\n",
        "dataset = {\n",
        "    'coords': data['coords'].reshape((-1, 2)),\n",
        "    'colors': data['colors'].reshape((-1, 3)),\n",
        "    'ids': data['ids'].reshape((-1, 1)),\n",
        "}\n",
        "\n",
        "num_items = np.prod(data['ids'].shape)\n",
        "perm = random.permutation(rng, jnp.arange(num_items))\n",
        "dataset = jax.tree_map(lambda x: x[perm], dataset)\n",
        "dataset = prepare_data(dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FCzdBUiRNspv"
      },
      "source": [
        "def compute_psnr(mse):\n",
        "  return -10. * jnp.log(mse) / jnp.log(10.)\n",
        "\n",
        "\n",
        "def loss_fn(params, batch, alpha):\n",
        "  color_pred = model.apply(params, batch['coords'], batch['ids'].squeeze(-1), alpha, True)['color']\n",
        "  loss = (color_pred - batch['colors']) ** 2\n",
        "  loss = loss.mean()\n",
        "  return loss\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def train_step(optimizer, batch, key, alpha, lr):\n",
        "  _, key = random.split(rng, 2)\n",
        "  grad_fn = jax.value_and_grad(loss_fn, argnums=0)\n",
        "  loss, grad = grad_fn(optimizer.target, batch, alpha)\n",
        "  loss = jax.lax.pmean(loss, axis_name='batch')\n",
        "  grad = jax.lax.pmean(grad, axis_name='batch')\n",
        "  new_optimizer = optimizer.apply_gradient(grad, learning_rate=lr)\n",
        "  return new_optimizer, key, loss\n",
        "\n",
        "\n",
        "def render_item(model, params, item):\n",
        "  target_image = item['colors']\n",
        "  pred = model.apply(\n",
        "      jax_utils.unreplicate(params), \n",
        "      item['coords'].reshape((-1, 2)),\n",
        "      item['ids'].reshape((-1,)),\n",
        "      alpha,\n",
        "      True)\n",
        "  pred_image = pred['color'].reshape(target_image.shape)\n",
        "  pred_flow = flow_to_image(np.array(pred['flow'].reshape((*image_shape, 2))))\n",
        "  pred_template = model.apply(\n",
        "      jax_utils.unreplicate(optimizer.target), \n",
        "      item['coords'].reshape((-1, 2)),\n",
        "      item['ids'].reshape((-1,)),\n",
        "      alpha,\n",
        "      False)['color'].reshape(target_image.shape)\n",
        "\n",
        "  return {\n",
        "      'target_color': target_image,\n",
        "      'pred_color': pred_image,\n",
        "      'pred_flow': pred_flow,\n",
        "      'pred_template': pred_template,\n",
        "  }\n",
        "\n",
        "\n",
        "# Train.\n",
        "max_iters = 10000\n",
        "# lr_schedule = schedules.ConstantSchedule(8e-3)\n",
        "lr_schedule = schedules.from_config({\n",
        "  'type': 'delayed',\n",
        "  'delay_steps': 50,\n",
        "  'delay_mult': 0.01,\n",
        "  'base_schedule': {\n",
        "    'type': 'exponential',\n",
        "    'initial_value': 8e-3,\n",
        "    'final_value': 8e-5,\n",
        "    'num_steps': max_iters,\n",
        "  },\n",
        "})\n",
        "if use_anneal:\n",
        "  alpha_schedule = schedules.LinearSchedule(\n",
        "      model.deform_min_freq_log2, model.deform_max_freq_log2+1, 5000)\n",
        "else:\n",
        "  alpha_schedule = schedules.ConstantSchedule(model.deform_max_freq_log2+1)\n",
        "\n",
        "\n",
        "optimizer_def = optim.Adam(lr_schedule(0))\n",
        "optimizer = optimizer_def.create(init_params)\n",
        "optimizer = jax_utils.replicate(optimizer, devices)\n",
        "\n",
        "p_train_step = jax.pmap(\n",
        "    train_step, axis_name='batch', devices=devices, in_axes=(0, 0, 0, None, None))\n",
        "\n",
        "tt = utils.TimeTracker()\n",
        "\n",
        "show_idx = 0\n",
        "keys = random.split(rng, len(devices))\n",
        "tt.tic('data')\n",
        "for i in range(max_iters):\n",
        "  batch = dataset\n",
        "  tt.toc('data')\n",
        "  if i > max_iters:\n",
        "    break\n",
        "  alpha = alpha_schedule(i)\n",
        "  lr = lr_schedule(i)\n",
        "  with tt.record_time('p_train_step'):\n",
        "    optimizer, keys, losses = p_train_step(optimizer, batch, keys, alpha, lr)\n",
        "  if i % 10 == 0:\n",
        "    loss = jax_utils.unreplicate(losses)\n",
        "    psnr = compute_psnr(loss)\n",
        "    print(f'{i}: lr = {lr:.04f}, alpha = {alpha:.02f}, loss = {loss:.04f}, psnr = {psnr:.02f}', end='\\r')\n",
        "\n",
        "  if i % 500 == 0:\n",
        "    print(f'{i}: loss = {loss:.04f}, psnr = {psnr:.02f}')\n",
        "    print(f'Showing image {show_idx}')\n",
        "    item = data_items[show_idx]\n",
        "    render = render_item(model, optimizer.target, item)\n",
        "    imedia.show_images([render['target_color'], \n",
        "                        render['pred_color'], \n",
        "                        jnp.abs(render['target_color'] - render['pred_color']), \n",
        "                        render['pred_template'], \n",
        "                        render['pred_flow']],\n",
        "                        titles=['GT', 'Pred', 'Abs. Error', 'Template', 'Flow'])\n",
        "    del render\n",
        "    show_idx = (show_idx + 1) % len(images)\n",
        "  tt.tic('data')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yBL8VVH4UPaP"
      },
      "source": [
        "batch['ids'].shape"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hR9E5jU8eFNA"
      },
      "source": [
        "def _save_image(path, image):\n",
        "  image = image_utils.image_to_uint8(image)\n",
        "  image_utils.save_image(path, image)\n",
        "\n",
        "exp_name = f'spaceman_{deform_type}_{\"annealed\" if use_anneal else \"notannealed\"}_{min_freq}_to_{max_freq}'\n",
        "save_dir = gpath.GPath('', exp_name)\n",
        "save_dir.mkdir(exist_ok=True, parents=True)\n",
        "\n",
        "\n",
        "for i, item in enumerate(data_items):\n",
        "  render = render_item(model, optimizer.target, item)\n",
        "  _save_image(save_dir / f'{i:02d}_target_color.png', render['target_color'])\n",
        "  _save_image(save_dir / f'{i:02d}_pred_color.png', render['pred_color'])\n",
        "  _save_image(save_dir / f'{i:02d}_pred_template.png', render['pred_template'])\n",
        "  _save_image(save_dir / f'{i:02d}_pred_flow.png', render['pred_flow'])\n",
        "  imedia.show_images([render['target_color'], \n",
        "                      render['pred_color'], \n",
        "                      jnp.abs(render['target_color'] - render['pred_color']), \n",
        "                      render['pred_template'], \n",
        "                      render['pred_flow']],\n",
        "                      titles=['GT', 'Pred', 'Abs. Error', 'Template', 'Flow'])\n",
        " "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FX4RS7bhTVmY"
      },
      "source": [
        "## Here be dragons."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qhhhSStyTZh9"
      },
      "source": [
        "0/0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_CbXe1t2j-AG"
      },
      "source": [
        "exp_dir = gpath.GPath('')\n",
        "exp_paths = sorted(x for x in exp_dir.iterdir() if x.name != 'training_images')\n",
        "image_id = 11\n",
        "\n",
        "target_color = image_utils.load_image(exp_paths[0] / f'{image_id:02d}_target_color.png')\n",
        "pred_color_paths = [p / f'{image_id:02d}_pred_color.png' for p in exp_paths]\n",
        "pred_colors = utils.parallel_map(image_utils.load_image, pred_color_paths, show_pbar=True)\n",
        "\n",
        "pred_template_paths = [p / f'{image_id:02d}_pred_template.png' for p in exp_paths]\n",
        "pred_templates = utils.parallel_map(image_utils.load_image, pred_template_paths, show_pbar=True)\n",
        "\n",
        "pred_flow_paths = [p / f'{image_id:02d}_pred_flow.png' for p in exp_paths]\n",
        "pred_flows = utils.parallel_map(image_utils.load_image, pred_flow_paths, show_pbar=True)\n",
        "\n",
        "imedia.show_images([target_color] + pred_colors, ['gt'] + [p.name for p in exp_paths])\n",
        "imedia.show_images([np.zeros_like(target_color)] + pred_templates, [''] + [p.name for p in exp_paths])\n",
        "imedia.show_images([np.zeros_like(target_color)] + pred_flows, [''] + [p.name for p in exp_paths])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1wKZCuPbSXcr"
      },
      "source": [
        "compute_psnr(0.049)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G10bhihNUjlm"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}